import os

import json5
import pandas as pd
import mysql.connector
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 设置字体为 SimHei
plt.rcParams['axes.unicode_minus'] = False

from docx import Document
from docx.shared import Inches
from macroeconomic_analysis_agent.utils import parse_code
from base_agent import BaseAgent

bar_prompt = """
你是一位数据科学家，你的目标是根据给定df中的列名与数据信息，给出生成柱状图的x轴列名与y轴列名，x轴列名使用时间相关列，
当前df信息:{df}
请严格按以下JSON格式回答，不得添加任何额外内容。:
```json
    {{
        "x": "时间列",
        "y": "列名1"
    }},
    ...
```

"""

data_scientist_prompt = """
你是一位数据科学家，你的目标是根据给定资源中的数据库结构信息，使用正确的MySQL语法的SQL语句来分析并解决用户问题。
请逐步思考以实现目标。同时，请严格遵守重要提醒中的约束和规范。

用户问题:
{user_input}

数据库类型:MySQL
MYSQL不支持LAG窗口函数
数据库结构信息:
表名:sina_data_2
中文表名:宏观经济分析数据源表
表注释: 宏观经济分析所需的数据，包含各类宏观经济相关数据
表类型: 增量表
表字段:
    字段名称:id; 字段类型:int(11); 字段描述:主键,
    字段名称: groupNameEn; 字段类型: varchar(255); 字段描述: 分组指标英文名称,
    字段名称: paramId; 字段类型: int(11); 字段描述: 暂无意义,
    字段名称: groupNameCn; 字段类型: varchar(255); 字段描述: 分组指标中文名称， 包含'季度GDP'，'货币供应量'，'CPI居民消费价格指数'， '按贸易分进出口总额'， '年度GDP'， '制造业PMI'
    字段名称: year; varchar(255); 字段描述: 年份,
    字段名称: norm; 字段类型: varchar(255); 字段描述: 详细指标名称 包含'国内生产总值同比增长'，'流通中现金(M0)', '流通中现金(M0)同比增长'， '货币(狭义货币M1)', '货币(狭义货币M1)同比增长', '货币和准货币（广义货币M2）', '货币和准货币（广义货币M2）同比增长', '价格指数'， '一般贸易出口'， '一般贸易进口'， '加工贸易出口'， '加工贸易进口'， '其他贸易出口'， '其他贸易进口'， '投资完成额'， '投资完成额同比增长'， '建筑业'， 'PMI'， '新订单指数'， '新出口订单指数'
    字段名称: data_value; 字段类型: varchar(255); 字段描述: 指标对应数值,
    字段名称: month; 字段类型: varchar(255); 字段描述: 月份或季度,
    字段名称: norm_conditon; 字段类型: varchar(255); 字段描述: 指标条件,
    字段名称: unit; 字段类型: varchar(255); 字段描述: 指标对应数值单位,

约束:
请使用简体中文
1. 请仔细检查生成的SQL语句。必须严格遵循提供的数据结构定义，严禁使用不存在的字段和数据值。不允许在表B中包含表A的字段。允许进行多表关联查询。此外，SQL语法必须严格遵守MySQL数据库语法规则；任何不被MySQL支持的语法、特性、函数或关键字都是严格禁止的。特别地，禁止使用WITH子句、FILTER子句、REAL数据类型以及CTE（公用表表达式）查询结构。
2. 请仅使用给定表结构中的字段构建查询语句，严禁为了构造此类查询而编造不存在的字段。
3. 生成的SQL语句应输出在一行内，不得包含任何换行字符。
4. 如果目标分析的数据和字段位于不同的表中，建议首先使用多表关联查询，并注意多个表结构之间的关联。
5. 禁止自行构造数据作为查询条件。如果要查询特定字段，如果提供了该字段的值，则可以对该字段进行分组统计查询。
6. 生成的sql应该只查询两个参数，一个是时间参数，一个是用户问题的目标值。
7. 如果出现了例子中的问题，请直接完整使用例子中的sql不要修改

注意：
近五年建筑业总产值增长率按月变化情况，日期为年不需要月份

请严格按以下JSON格式回答，不得添加任何额外内容。:
```json
    {{
        "sql": "Executable sql generated for the current target/problem",
        "thought": "Summary of thoughts to the user",
        "title": "Summary of User's Question Title, Please Output in Chinese (Title Should Not Include "Query" or Similar Words)"
    }},
    ...
```
确保响应是正确的JSON格式，并且可以被Python的json.loads解析。

例子1:
查询并分析近二十年GDP增长率按季度变化情况的数据
sql: 
SELECT CONCAT(`year`, '-', LPAD(`month`, 2, '0')) AS '日期', `data_value` as '国内生产总值同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%国内生产总值同比增长%' AND `groupNameCn` LIKE '%季度GDP%' AND `year` >= '2003' order by CAST(`year` AS UNSIGNED) ASC, CAST(`month` AS UNSIGNED) ASC约120条记录

例子2:
查询并分析近五年M0增速按月变化情况的数据
sql:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm LIKE '%流通中现金(M0)%' AND norm NOT LIKE '%同比增长%' THEN data_value END) AS '流通中现金(M0)',MAX(CASE WHEN norm LIKE '%流通中现金(M0)同比增长%' THEN data_value END) AS '流通中现金(M0)同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%M0%' AND `groupNameCn` LIKE '%货币供应量%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

例子3:
查询并分析近五年M1增速按月变化情况的数据
sql:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm LIKE '%货币(狭义货币M1)%' AND norm NOT LIKE '%同比增长%' THEN data_value END) AS '货币(狭义货币M1)',MAX(CASE WHEN norm LIKE '%货币(狭义货币M1)同比增长%' THEN data_value END) AS '货币(狭义货币M1)同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%M1%' AND `groupNameCn` LIKE '%货币供应量%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

例子4:
查询并分析近五年M2增速按月变化情况的数据
sql:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm LIKE '%货币和准货币（广义货币M2）%' AND norm NOT LIKE '%同比增长%' THEN data_value END) AS '货币和准货币（广义货币M2）',MAX(CASE WHEN norm LIKE '%货币和准货币（广义货币M2）同比增长%' THEN data_value END) AS '货币和准货币（广义货币M2）同比增长' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%M2%' AND `groupNameCn` LIKE '%货币供应量%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

例子5:
查询并分析近十年CPI指数按月变化情况的数据
sql:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期', `data_value` as 'CPI居民消费价格指数' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%价格指数%' AND `groupNameCn` LIKE '%CPI居民消费价格指数%' AND `year` >= '2015' ORDER BY CAST(`year` AS UNSIGNED) ASC, CAST(`month` AS UNSIGNED) ASC;

例子6:
查询并分析2015年来进出口总额数据及变化情况，包括进口与出口产品分类数据
sql:
SELECT YEAR AS '日期',MAX(CASE WHEN norm LIKE '%一般贸易出口%' THEN data_value ELSE NULL END) AS '一般贸易出口的值',MAX(CASE WHEN norm LIKE '%一般贸易进口%' THEN data_value ELSE NULL END) AS '一般贸易进口',MAX(CASE WHEN norm LIKE '%加工贸易出口%' THEN data_value ELSE NULL END) AS '加工贸易出口',MAX(CASE WHEN norm LIKE '%加工贸易进口%' THEN data_value ELSE NULL END) AS '加工贸易进口',MAX(CASE WHEN norm LIKE '%其他贸易出口%' THEN data_value ELSE NULL END) AS '其他贸易出口',MAX(CASE WHEN norm LIKE '%其他贸易进口%' THEN data_value ELSE NULL END) AS '其他贸易进口' FROM `data`.`sina_data_2` WHERE `groupNameCn` LIKE '%按贸易分进出口总额%' AND `year`>='2015' GROUP BY YEAR ORDER BY CAST(YEAR AS UNSIGNED) ASC;

例子7:
查询并分析近五年投资完成额按月变化情况的数据
sql:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm='投资完成额' THEN data_value ELSE NULL END) AS '投资完成额的值',MAX(CASE WHEN norm='投资完成额同比增长' THEN data_value ELSE NULL END) AS '投资完成额同比增长的值' FROM `data`.`sina_data_2` WHERE `norm` IN ('投资完成额','投资完成额同比增长') AND `groupNameCn` LIKE '%固定资产投资%' AND `year`>='2020' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;

例子8:
查询并分析近五年建筑业总产值增长率按月变化情况的数据
sql:
SELECT `YEAR` AS '日期', `data_value` as '建筑业总产值' FROM `data`.`sina_data_2` WHERE `norm` LIKE '%建筑业%' AND `year` >= '2020' AND `groupNameCn` = '年度GDP' ORDER BY CAST(`year` AS UNSIGNED) ASC;

例子9:
查询并分析近五年PMI指数、新出口订单指数、新订单指数按月变化情况的数据
sql:
SELECT CONCAT(YEAR,'-',LPAD(MONTH,2,'0')) AS '日期',MAX(CASE WHEN norm='PMI' THEN data_value ELSE NULL END) AS 'PMI',MAX(CASE WHEN norm='新订单指数' THEN data_value ELSE NULL END) AS '新订单指数',MAX(CASE WHEN norm='新出口订单指数' THEN data_value ELSE NULL END) AS '新出口订单指数' FROM `data`.`sina_data_2` WHERE `norm` IN ('PMI','新订单指数','新出口订单指数') AND `year`>='2020' AND `groupNameCn` LIKE '%制造业PMI%' GROUP BY YEAR,MONTH ORDER BY CAST(YEAR AS UNSIGNED) ASC,CAST(MONTH AS UNSIGNED) ASC;
"""

analysis_prompt = """
你是一个宏观经济分析师，你的目标是根据问题标题，以及对应数据进行宏观经济相关的分析，并输出分析内容，内容包含数据以及对数据的解读控制在250-300字。
问题标题: {title}
对应数据: {data}
"""


class DataScientistAgent(BaseAgent):

    def __init__(self, task, model, plan_id, **kwargs):
        super().__init__(model=model, **kwargs)  # 调用基类构造方法
        # 以下为该类特有的属性初始化
        self.plan_id = plan_id  # 计划ID，用于标识或追踪计划
        self.user_input = task['instruction']  # 用户的输入，来源于任务的指令部分
        # 根据用户输入格式化数据科学家的提示信息
        self.data_scientist_prompt = data_scientist_prompt.format(user_input=self.user_input)
        self.img_path = []

    def generate_sql(self):
        """
        生成SQL语句的方法。

        该方法通过调用语言模型来生成SQL语句。它会尝试多次调用语言模型，
        直到成功获取到有效的SQL语句或达到最大尝试次数。语言模型的响应
        会被解析并转换为JSON对象，从而提取出所需的SQL语句和标题。
        """
        # 初始化语言模型调用成功标志为False
        call_llm_success = False
        # 初始化语言模型调用次数计数器
        call_llm_count = 0
        # 初始化任务文本为空字符串
        tasks_text = ''
        # 构建发送给语言模型的消息，当前只有一个用户角色的消息
        messages = [{'role': 'user', 'content': self.data_scientist_prompt}]

        # 循环尝试调用语言模型，直到成功或达到最大尝试次数（10次）
        while not call_llm_success and call_llm_count < 10:
            # 调用语言模型并获取响应
            resp = self._call_llm(messages=messages)
            # 重置任务文本
            tasks_text = ''
            # 遍历响应内容并累加到任务文本中
            for r in resp:
                tasks_text += r
            # 检查任务文本中是否包含错误代码
            if 'Error code' in tasks_text:
                # 如果包含错误代码，则增加调用尝试次数
                call_llm_count += 1
            else:
                # 如果不包含错误代码，则设置成功标志为True，表示调用成功
                call_llm_success = True

        # 打印任务文本
        print('Tasks_text: ', tasks_text)
        # 将任务文本解析为指定语言的代码（此处为JSON）
        tasks_text = parse_code(text=tasks_text, lang='json')
        # 将解析后的任务文本解析为JSON对象
        tasks = json5.loads(tasks_text)
        # 从解析后的JSON对象中提取SQL语句和标题，并赋值给类的属性
        self.sql = tasks['sql']
        self.title = tasks['title']

    def execute_sql(self):
        """
        执行SQL查询并获取结果数据框。

        该方法连接到指定的数据库，执行实例对象中的SQL语句，并将结果存储在数据框中。
        若执行过程中出现异常，则被捕获，避免程序崩溃。
        """
        try:
            # 建立与数据库的连接
            connection = mysql.connector.connect(
                host='81.68.216.231',  # 数据库服务器地址
                database='data',  # 数据库名
                user='root',  # 用户名
                password='Wb@980324'  # 密码
            )
            print(f"当前生成的sql：{self.sql}")
            # 使用SQL语句读取数据库数据，并存储在数据框中
            df = pd.read_sql(self.sql, connection)
            self.df = df  # 将查询结果赋值给实例变量df
        except Exception as e:
            # 捕获并处理异常，避免程序因数据库错误而崩溃
            print(e)

    def generate_bar_chart(self):
        """
        生成柱状图。

        本方法通过调用语言模型生成柱状图的相关信息，然后利用这些信息绘制柱状图并保存。
        """
        # 准备绘制柱状图的提示信息，包含数据帧的前3行
        df_prompt = bar_prompt.format(df=str(self.df.head(3)))
        # 初始化消息列表，用于与语言模型交互
        messages = [{'role': 'user', 'content': df_prompt}]
        # 调用语言模型，获取绘制柱状图所需的信息
        resp = self._call_llm(messages=messages)
        # 解析语言模型的响应，转换为JSON格式
        bar_info = parse_code(text=resp, lang='json')
        # 将解析的JSON信息转换为字典
        bar_info = json5.loads(bar_info)
        # 将数据帧中的y轴数据转换为数值类型
        self.df[bar_info['y']] = pd.to_numeric(self.df[bar_info['y']], errors='coerce')
        # 使用数据帧中的x和y数据绘制柱状图
        self.df.plot(x=bar_info['x'], y=bar_info['y'], kind='bar', title=self.title)
        # 设置图表标签
        plt.xlabel(bar_info['x'])
        plt.ylabel(bar_info['y'])
        # 准备保存图表的文件路径
        folder_path = f'./{self.plan_id}/'  # 当前目录下的 images 文件夹
        # 构造文件完整路径
        file_path = os.path.join(folder_path, f'{self.title}.png')
        # 保存图表到本地文件
        plt.savefig(file_path)
        # 记录图表文件路径
        self.img_path = file_path

    def generate_line_chart(self):
        """
        生成折线图

        本函数用于根据保存的数据生成折线图，并将图表保存为图片文件。
        先判断数据中是否存在空值，如果有则用0填充。
        然后获取除日期外的所有列名，针对每一列数据绘制一个折线图。
        图表的标题、x轴标签、y轴标签均根据列名设置。
        通过调整x轴标签的旋转角度和刻度数量来优化图表的可读性。
        最后将每个绘制的图表保存为png文件，并记录每个图表的文件路径。
        """
        # 复制数据集以避免修改原始数据
        df = self.df.copy()

        # 检查并处理数据集中的空值
        if df.isnull().values.any():
            df.fillna(0, inplace=True)
        # 获取除日期以外的列名
        value_columns = df.columns.drop('日期')

        # 对于每列生成一个折线图
        for column in value_columns:
            df[column] = pd.to_numeric(df[column], errors='coerce')
            # 将 NaN 替换为 0
            df[column].fillna(0, inplace=True)

            plt.figure(figsize=(10, 6))

            # 绘制折线图
            plt.plot(df['日期'], df[column], marker='o', label=column)

            # 设置标题和坐标轴标签
            plt.title(column)
            plt.xlabel('日期')
            plt.ylabel(column)

            # 自动旋转x轴标签以防止重叠
            plt.xticks(rotation=45)

            # 获取y轴的最大值和最小值
            y_min = df[column].min()
            y_max = df[column].max()

            # 设置x轴刻度间隔
            ax = plt.gca()  # 获取当前axes实例
            nbins = df.shape[0] if df.shape[0] < 12 else 12
            ax.xaxis.set_major_locator(MaxNLocator(nbins=nbins, prune='both'))  # 可以调整nbins的值来改变显示的标签数量

            # 设置y轴刻度，最多显示12个刻度
            # 计算y轴的刻度值
            y_range = y_max - y_min
            if y_range > 0:
                step_size = y_range / 11  # 11是因为我们要包括最小值和最大值在内的12个刻度
                y_ticks = np.arange(y_min, y_max + step_size, step_size)
            else:
                y_ticks = [y_min]

            # 设置y轴的刻度
            ax.set_yticks(y_ticks)

            # 显示网格
            plt.grid(True)

            # 图例
            plt.legend()

            # 准备保存图表的文件路径
            folder_path = f'./{self.plan_id}/'  # 当前目录下的 images 文件夹
            # 构造文件完整路径
            file_path = os.path.join(folder_path, f'{column.replace(" ", "_")}.png')
            # 保存图表到本地文件
            plt.savefig(file_path)
            # 记录图表文件路径
            self.img_path.append(file_path)

    def generate_analysis_content(self):
        """
        生成分析内容。

        根据分析的标题和数据，构造一个分析提示（ana_prompt），并通过调用语言模型（LLM）获得分析内容。
        这个过程首先格式化一个分析提示字符串，该字符串包含分析的标题和数据信息。然后，创建一个消息列表，
        其中包含一个以用户角色发送的分析提示。接下来，调用_llm方法传递该消息列表以获取响应。最后，将
        获取到的分析内容赋值给self.analysis_content属性。
        """
        # 格式化分析提示字符串，包含分析标题和数据信息
        ana_prompt = analysis_prompt.format(title=self.title, data=self.df.to_string())
        print(f"问题： {ana_prompt}")
        # 创建一个消息列表，包含一个以用户角色发送的分析提示
        messages = [{'role': 'user', 'content': ana_prompt}]

        # 调用语言模型获取分析内容
        resp = self._call_llm(messages=messages)

        # 将获取到的分析内容赋值给self.analysis_content属性
        self.analysis_content = resp

    def update_or_create_document(self):
        """
        更新或创建文档。

        根据计划ID创建一个Word文档。如果文档已经存在，则加载现有文档；
        否则创建一个新的文档。然后将分析内容作为图片和文本添加到文档中，
        并保存文档。
        """
        # 生成文件名，使用计划ID作为前缀
        file_name = self.plan_id + "_document.docx"
        # 构建文件的完整路径
        file_path = os.path.join(f'./{self.plan_id}/', file_name)

        # 检查文档是否已经存在
        if not os.path.exists(file_path):
            # 如果文件不存在，创建一个新的文档
            doc = Document()
        else:
            # 如果文件存在，加载现有的文档
            doc = Document(file_path)

        # 向文档中添加图片，设置图片宽度
        for img_path in self.img_path:
            doc.add_picture(img_path, width=Inches(5.5))
        # 向文档中添加分析内容的段落
        doc.add_paragraph(self.analysis_content)
        # 保存文档
        doc.save(file_path)

    def run(self):
        """
        运行数据分析流程。

        该方法按顺序执行一系列数据分析步骤，包括生成SQL语句、执行SQL语句、
        生成柱状图、生成分析内容以及更新或创建文档。如果整个流程无异常，
        则返回True，否则捕获异常并打印错误信息，返回False。

        Returns:
            bool: 如果整个流程成功则返回True，否则返回False。
        """
        try:
            # 生成SQL语句
            self.generate_sql()

            # 执行SQL语句
            self.execute_sql()

            # # 生成柱状图
            # self.generate_bar_chart()

            # 生成折线图
            self.generate_line_chart()

            # 生成分析内容
            self.generate_analysis_content()

            # 更新或创建文档
            self.update_or_create_document()

            return True
        except Exception as e:
            # 打印异常信息
            print(e)
            return False


if __name__ == '__main__':
    d = DataScientistAgent(task={'instruction': '查询并分析2015年至今的中国基于购买力平价的国内生产总值在世界总量中的份额'}, model='qwen2-72b-instruct')
    d.run()